import math
import traceback
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QAction
from PyQt5.QtChart import QChart, QLineSeries, QValueAxis
from UI1 import *
from main import mainAllocation
import pickle
from const import *


class MyMainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        super(MyMainWindow, self).__init__(parent)

        # =======================运算区===================
        self.mainallocation = mainAllocation()

        # ================显示区==================
        self.setupUi(self)
        # 绘图
        self._plotreservoirview()
        #
        self.statusbar.showMessage('正常启动', 5*1000)

        #
        self.reservoirTypeBox.addItems(['水库', '河流'])
        self.reservoirviewcombo.addItems(['水位库容曲线', '水位过程', '库容过程'])
        self.constructTypeBox.addItems(['单向闸', '双向闸', '泵站'])
        self.constructviewcombo.addItems(['流量过程'])


        # self.currentreservoir_object = []    # 当前选中的储水池对象
        # self.currentconstruct_object = []    # 当前选中的建筑物对象
        self._showGlobe()


    # 测试集
    def testbutton(self):
        # print('按钮被点击了一次')
        # filename = QFileDialog.getOpenFileName(self, '选取文件','','')
        filename = QFileDialog.getSaveFileName(self, '保存文件', 'test.pickle', '.pickle')
        print(filename[0])

    def testedit(self, arg__1):
        # 获得对话框中的文字（只要有变动均会获取）
        # print('当前对话框的文字为',arg__1)
        pass

    def testbox(self, int__1):
        # 获得下拉菜单中所选取的序号，注意仅当选择发生变化时才会产生信号
        # print('当前下拉菜单选择的是%s++'%int__1)
        pass


    def restrict_values(values):
        # 装饰器
        # 上述代码中，我们定义了一个名为restrict_values的装饰器函数，它接受一个参数values，即允许的值的列表。
        # 装饰器函数内部又定义了一个名为decorator的装饰器函数，它接受一个参数func，即被装饰的函数或方法。
        # decorator函数内部又定义了一个名为wrapper的嵌套函数，它实际执行函数或方法的逻辑。
        # 在wrapper函数内部，我们首先检查传入参数的值是否在允许的值列表中，如果不在则抛出异常。
        # 然后，我们调用被装饰的函数或方法，并返回其返回值。
        # 使用装饰器时，我们只需要在定义函数或方法的时候，在其上面添加@restrict_values(['Alice', 'Bob'])
        # 即可实现限定传入参数为指定的值。
        def decorator(func):
            def wrapper(*args, **kwargs):
                # print(args, values)
                if args[1] not in values:
                    raise ValueError('Invalid value')
                return func(*args, **kwargs)
            return wrapper
        return decorator

    # ============================================正式开始=================================
    # =====================菜单栏
    def getmainallocation(self):
        # 获取已有的数据
        self.statusbar.showMessage("读取已有数据", 3 * 1000)
        try:
            filename = QFileDialog.getOpenFileName(self, '打开文件', 'test.pickle','')
            if filename[0]:
                self.statusbar.showMessage("打开中......")
                with open(filename[0], 'rb') as file:
                    self.mainallocation = pickle.load(file)
                    self.statusbar.showMessage("打开成功")
            else:
                self.statusbar.showMessage("取消打开")
        except Exception as e:
            self.statusbar.showMessage("打开不成功,错误代码：" + str(e))
            print(str(e))
        pass

    def save(self):
        self.statusbar.showMessage("保存数据为python格式", 3*1000)
        try:
            filename = QFileDialog.getSaveFileName(self, '保存文件', 'test.pickle', '.pickle')
            if filename[0]:
                self.statusbar.showMessage("保存中......")
                with open(filename[0], 'wb') as file:
                    pickle.dump(self.mainallocation, file)
                self.statusbar.showMessage("保存成功")
            else:
                self.statusbar.showMessage("取消保存")
        except Exception as e:
            self.statusbar.showMessage("保存不成功,错误代码：" + str(e))
            print(str(e))
        pass
    def out(self):
        self.statusbar.showMessage("退出", 3*1000)
        pass

    def example(self):
        self.statusbar.showMessage("读取样例数据", 3 * 1000)
        try:
            with open('test.pickle', 'rb') as file:
                self.mainallocation = pickle.load(file)
                # print('读取成功')
                # print(self.mainallocation.getreservoirname())

                self._recurrent('ALL')

                self._showAll()
        except Exception as e:
            self.statusbar.showMessage("打开不成功,错误代码：" + str(e))
            print(str(e))
        pass

    def help(self):
        pass

    # ===========全局设置
    def resetconst(self):
        self.mainallocation.Step = int(self.stepEdit.text())
        self.mainallocation.Steptime = float(self.steptimeEdit.text())  # 每一步的计算步长，单位是秒
        self.mainallocation.Relevel = float(self.relevelEdit.text())

    def reflash(self):
        self.mainallocation.reflash()
        self._recurrent('ALL')
        self.statusbar.showMessage("已重置计算序列", 3 * 1000)
        pass

    def allocation(self):
        self.statusbar.showMessage("运算ing", 3 * 1000)
        try:
            self.mainallocation.allocation()
            self._recurrent('ALL')
            self.statusbar.showMessage("运算成功", 3 * 1000)
        except Exception as e:
            print('未能成功运算', traceback.print_exc(), e)
            self.statusbar.showMessage('失败，'+ str(e),traceback.print_exc(), 3 * 1000)

    # ===========储水池=====
    def addreservoir(self):
        self.statusbar.showMessage("增加储水池", 3 * 1000)
        try:
            exception = self.mainallocation.addreservoir(
                serial=float(self.reservoirSerialEdit.text()),
                name=self.reservoirNameEdit.text(),
                firstlevel=float(self.reservoirFirstLevelEdit.text()),
                upconstruct_serial=self.strToListNum(self.upconstructserialEdit.text()),
                downconstruct_serial=self.strToListNum(self.downconstructserialEdit.text()),
                reservoirtype='reservoir' if self.reservoirTypeBox.currentIndex() == 0 else 'river',
                levelvolumerelation=[self.strToListNum(self.levelrelationEdit.text()), self.strToListNum(self.volumerelationEdit.text())]
            )
            if exception:
                self.statusbar.showMessage(str(exception), 3 * 1000)
            else:
                self.currentreservoir_object = self.mainallocation.getreservoir(float(self.reservoirSerialEdit.text()))
                self._showreservoir()
        except Exception as e:
            print('未能成功添加', traceback.print_exc(), e)
            self.statusbar.showMessage('失败，'+ str(e),traceback.print_exc(), 3 * 1000)
            return e


    def currentreservoir(self, int__1):
        # 获得下拉菜单中所选取的序号，注意仅当选择发生变化时才会产生信号
        # self.statusbar.showMessage("当前选择%s号储水池" %int__1, 3 * 1000)
        # print(type(int__1))
        self.currentreservoir_object = self.mainallocation.reservoirlist[int__1]
        self.statusbar.showMessage("当前选择%s号储水池" % self.currentreservoir_object.serial, 3 * 1000)
        self._showreservoir()
        pass

    def delreservoir(self):
        # 删除当前下拉菜单选中的储水池
        self.mainallocation.reservoirlist.remove(self.currentreservoir_object)    # 删除当前值
        self._recurrent('reservoir')
        self._showreservoir()
        self.statusbar.showMessage("删除储水池", 3 * 1000)
        pass

    def rereservoir(self):
        # 删除当前，再重新添加
        try:
            self.mainallocation.reservoirlist.remove(self.currentreservoir_object)  # 删除当前值
            self.addreservoir()
            self.statusbar.showMessage("储水池更新成功", 3 * 1000)
        except Exception as e:
            print('未能成功添加', traceback.print_exc(), e)
            self.statusbar.showMessage('失败，' + str(e), traceback.print_exc(), 3 * 1000)
            return e

    def reservoirview(self):
        # self.statusbar.showMessage("选中", 3 * 1000)
        if self.currentreservoir_object:
            self._showreservoirview()
        # '水位库容曲线', '水位过程', '库容过程'
        pass

    # ===========建筑物

    def addconstruct(self):
        self.statusbar.showMessage("增加建筑物成功", 3 * 1000)
        try:
            exception = self.mainallocation.addconstruct(
                serial=float(self.constructSerialEdit.text()),
                name=self.constructNameEdit.text(),
                constructtype= self.constructTypeBox.currentText(),
                upserial=float(self.upreservoirserialEdit.text()),
                downserial=float(self.downreservoirserialEdit.text()),
                baselevel=float(self.baselevelEdit.text()),
                width=float(self.widthEdit.text()),
                hight=float(self.hightEdit.text()),
                maxflow=float(self.maxflowEdit.text()),
                rule=self.strToListNum(self.ruleEdit.text())
            )
            if exception:
                self.statusbar.showMessage(str(exception), 3 * 1000)
            else:
                self.currentconstruct_object = self.mainallocation.getconstruct(float(self.constructSerialEdit.text()))
                self._showconstruct()
        except Exception as e:
            print('未能成功添加', traceback.print_exc(), e)
            self.statusbar.showMessage('失败，' + str(e), traceback.print_exc(), 3 * 1000)

    def currentconstruct(self, int__1:int):
        # 获得下拉菜单中所选取的序号，注意仅当选择发生变化时才会产生信号
        self.currentconstruct_object = self.mainallocation.constructlist[int__1]
        self.statusbar.showMessage("当前选择%s号建筑物" % self.currentconstruct_object.serial, 3 * 1000)
        self._showconstruct()
        pass

    def delconstruct(self):
        # 删除当前下拉菜单选中的建筑物
        self.mainallocation.constructlist.remove(self.currentconstruct_object)    # 删除当前值
        self._recurrent('construct')
        self._showconstruct()
        self.statusbar.showMessage("删除建筑物", 3 * 1000)
        pass

    def reconstruct(self):
        # 删除当前，再重新添加
        try:
            self.mainallocation.constructlist.remove(self.currentconstruct_object)  # 删除当前值
            self.addconstruct()
            self.statusbar.showMessage("建筑物更新成功", 3 * 1000)
        except Exception as e:
            print('未能成功添加', traceback.print_exc(), e)
            self.statusbar.showMessage('失败，' + str(e), traceback.print_exc(), 3 * 1000)
            return e

    # def rule(self):
    #     self.statusbar.showMessage("建筑物规则", 3 * 1000)
    #     pass

    def constructview(self, int__1):
        self.statusbar.showMessage("选择图", 3 * 1000)
        self._showconstructview()
        pass

    # ================== 内置函数 =================
    def _showAll(self):
        self._showGlobe()
        self._showreservoir()
        self._showconstruct()

    def _showGlobe(self):
        # 显示
        try:
            self.steptimeEdit.setText(str(self.mainallocation.Steptime))
            self.stepEdit.setText(str(self.mainallocation.Step))
            self.relevelEdit.setText(str(self.mainallocation.Relevel))  # 水位要求
        except Exception as e:
            print(traceback.print_exc(),e)


    def _showreservoir(self):
        # 蓄水池部分显示刷新
        self.reservoirlabel.setText('现有%s个储水池' % len(self.mainallocation.reservoirlist))
        # 下拉菜单增加几个按键
        self.reservoirlistcomboBox.clear()
        self.reservoirlistcomboBox.addItems(str(serial) for serial in self.mainallocation.getreservoirserial())

        if self.currentreservoir_object:
            self.reservoirlistcomboBox.setCurrentIndex(self.mainallocation.reservoirlist.index(self.currentreservoir_object))
            self.reservoirNameEdit.setText(str(self.currentreservoir_object.name))    # 名称
            self.reservoirSerialEdit.setText(str(self.currentreservoir_object.serial))    # 序号
            if self.currentreservoir_object.reservoirtype == 'reservoir':    # 设置类型
                self.reservoirTypeBox.setCurrentIndex(0)
            else:
                self.reservoirTypeBox.setCurrentIndex(1)
            self.reservoirFirstLevelEdit.setText(str(self.currentreservoir_object.waterlevellist[0])) # 初始水位
            self.upconstructserialEdit.setText(','.join([str(x) for x in self.currentreservoir_object.upconstruct_serial])) # 上游建筑物序号
            self.downconstructserialEdit.setText(','.join([str(x) for x in self.currentreservoir_object.downconstruct_serial])) # 下游建筑物序号

            self.levelrelationEdit.setText(','.join([str(x) for x in self.currentreservoir_object.levelvolumerelation[0]]))
            self.volumerelationEdit.setText(','.join([str(x) for x in self.currentreservoir_object.levelvolumerelation[1]]))
            self.reservoirviewcombo.setCurrentIndex(0)
            self._showreservoirview()

    def _showconstruct(self):
        # 建筑物部分展示刷新
        self.constructlabel.setText('现有%s个建筑物' % len(self.mainallocation.constructlist))
        # 下拉菜单增加几个按键
        self.constructlistcomboBox.clear()
        self.constructlistcomboBox.addItems(str(serial) for serial in self.mainallocation.getconstructserial())

        try:
            if self.currentconstruct_object:
                self.constructlistcomboBox.setCurrentIndex(
                    self.mainallocation.constructlist.index(self.currentconstruct_object))
                self.constructNameEdit.setText(str(self.currentconstruct_object.name))  # 名称
                self.constructSerialEdit.setText(str(self.currentconstruct_object.serial))  # 序号
                if self.currentconstruct_object.constructtype == '单向闸':  # 设置类型
                    self.constructTypeBox.setCurrentIndex(0)
                elif self.currentconstruct_object.constructtype == '双向闸':  # 设置类型
                    self.constructTypeBox.setCurrentIndex(1)
                elif self.currentconstruct_object.constructtype == '泵站':  # 设置类型
                    self.constructTypeBox.setCurrentIndex(2)
                self.upreservoirserialEdit.setText(str(self.currentconstruct_object.upserial))  # 上游储水池序号
                self.downreservoirserialEdit.setText(str(self.currentconstruct_object.downserial))  # 下游储水池序号
                self.baselevelEdit.setText(str(self.currentconstruct_object.baselevel))  # 底板高层
                self.maxflowEdit.setText(str(self.currentconstruct_object.maxflow))  # 最大流量
                self.widthEdit.setText(str(self.currentconstruct_object.width))  # 宽
                self.hightEdit.setText(str(self.currentconstruct_object.hight))  # 高
                self.constructviewcombo.setCurrentIndex(0)  # 图表
                self._showconstructview()  # 显示图表
        except Exception as e:
            print('', traceback.print_exc(), e)
            self.statusbar.showMessage('失败，' + str(e), traceback.print_exc(), 3 * 1000)

    def _showreservoirview(self):
        # ['水位库容曲线', '水位过程', '库容过程']
        chart_view = self.reservoirView
        if self.reservoirviewcombo.currentIndex() == 0:
            # print('当前选中水位库容曲线')
            chart = self._watervolumechart()
        elif self.reservoirviewcombo.currentIndex() == 1:
            # print('当前选中水位过程')
            chart = self._waterlistchart()
        elif self.reservoirviewcombo.currentIndex() == 2:
            # print('当前选中库容过程')
            chart = self._volumelistchart()
        chart_view.setChart(chart)
        pass

    def _showconstructview(self):
        # ['流量过程']
        chart_view = self.constructView
        chart = self._flowlistchart()
        chart_view.setChart(chart)
        pass

    def _watervolumechart(self):
        # 绘制水位库容曲线
        levellist = self.currentreservoir_object.levelvolumerelation[0]
        volumelist = self.currentreservoir_object.levelvolumerelation[1]

        chart = QChart()
        chart.legend().setVisible(True)
        chart.legend().setAlignment(QtCore.Qt.AlignBottom)

        series = QLineSeries()
        series.setName('水位库容曲线')    # 线名称
        for i in range(len(levellist)):
            series.append(volumelist[i], levellist[i])
        # 将QLineSeries对象添加到QChart对象中
        chart.addSeries(series)
        # 设置图表的横坐标和纵坐标
        axis_x = QValueAxis()
        axis_x.setTitleText('库容(立方米)')
        axis_y = QValueAxis()
        axis_y.setTitleText('水位(米)')
        chart.setAxisX(axis_x, series)
        chart.setAxisY(axis_y, series)

        return chart

    def _waterlistchart(self):
        # 绘制水位过程曲线
        levellist = self.currentreservoir_object.waterlevellist

        chart = QChart()
        chart.legend().setVisible(True)
        chart.legend().setAlignment(QtCore.Qt.AlignBottom)

        series = QLineSeries()
        series.setName('水位过程曲线')    # 线名称
        for i in range(len(levellist)):
            series.append(i, levellist[i])
        # 将QLineSeries对象添加到QChart对象中
        chart.addSeries(series)
        # 设置图表的横坐标和纵坐标
        axis_x = QValueAxis()
        axis_x.setTitleText('时段')
        axis_y = QValueAxis()
        axis_y.setTitleText('水位(米)')
        chart.setAxisX(axis_x, series)
        chart.setAxisY(axis_y, series)
        return chart

    def _volumelistchart(self):
        # 绘制库容过程曲线
        volumelist = self.currentreservoir_object.watervolumelist

        chart = QChart()
        chart.legend().setVisible(True)
        chart.legend().setAlignment(QtCore.Qt.AlignBottom)

        series = QLineSeries()
        series.setName('库容过程曲线')    # 线名称
        for i in range(len(volumelist)):
            series.append(i, volumelist[i])
        # 将QLineSeries对象添加到QChart对象中
        chart.addSeries(series)
        # 设置图表的横坐标和纵坐标
        axis_x = QValueAxis()
        axis_x.setTitleText('时段')
        axis_y = QValueAxis()
        axis_y.setTitleText('库容(立方米)')
        chart.setAxisX(axis_x, series)
        chart.setAxisY(axis_y, series)
        return chart

    def _flowlistchart(self):
        # 绘制过流过程曲线
        flowlist = self.currentconstruct_object.flowlist

        chart = QChart()
        chart.legend().setVisible(True)
        chart.legend().setAlignment(QtCore.Qt.AlignBottom)

        series = QLineSeries()
        series.setName('建筑物过流曲线')    # 线名称
        for i in range(len(flowlist)):
            series.append(i, flowlist[i])
        # 将QLineSeries对象添加到QChart对象中
        chart.addSeries(series)
        # 设置图表的横坐标和纵坐标
        axis_x = QValueAxis()
        axis_x.setTitleText('时段')
        axis_y = QValueAxis()
        axis_y.setTitleText('流量')
        chart.setAxisX(axis_x, series)
        chart.setAxisY(axis_y, series)
        return chart

    def _plotreservoirview(self, view:QChartView = []):
        # 更改view
        if view:
            chart_view = view
        else:
            chart_view = self.reservoirView
        # 创建QChart对象并设置属性
        levelvolumerelation0 = [[0,5,20,100],[0,50,200,1000]]    # levellist,volumelist
        levellist = levelvolumerelation0[0]
        volumelist = levelvolumerelation0[1]

        levellist = [1]
        volumelist = []

        chart = QChart()
        # chart.setTitle('My Chart')    # 无需标题
        chart.legend().setVisible(True)
        chart.legend().setAlignment(QtCore.Qt.AlignBottom)

        series = QLineSeries()
        series.setName('水位库容曲线')    # 线名称
        for i in range(len(levellist)):
            series.append(i, levellist[i])
        # 将QLineSeries对象添加到QChart对象中
        chart.addSeries(series)
        # 设置图表的横坐标和纵坐标
        axis_x = QValueAxis()
        axis_x.setTitleText('库容')
        axis_y = QValueAxis()
        chart.setAxisX(axis_x, series)
        chart.setAxisY(axis_y, series)

        # 将QChart对象设置给QChartView控件
        chart_view.setChart(chart)
        pass

    @restrict_values(['ALL', 'reservoir','construct'])
    def _recurrent(self, type):
        # 重置当前选中
        if type == 'ALL':
            if self.mainallocation.reservoirlist:
                self.currentreservoir_object = self.mainallocation.reservoirlist[0]  # 当前选中的储水池对象
            else:
                self.currentreservoir_object = []
            if self.mainallocation.constructlist:
                self.currentconstruct_object = self.mainallocation.constructlist[0]  # 当前选中的建筑物对象
            else:
                self.currentconstruct_object = []
        elif type == 'reservoir':
            if self.mainallocation.reservoirlist:
                self.currentreservoir_object = self.mainallocation.reservoirlist[0]  # 当前选中的储水池对象
            else:
                self.currentreservoir_object = []
        elif type == 'construct':
            if self.mainallocation.constructlist:
                self.currentconstruct_object = self.mainallocation.constructlist[0]  # 当前选中的建筑物对象
            else:
                self.currentconstruct_object = []
        self._showAll()

    # ===============其他转换功能===========
    def strToListNum(self, text:str):
        # print('text', text)
        if text:
            listNum = text.split(',')
            # print('listNum', listNum)
            listNum = [float(x) for x in listNum]
            # print(listNum)
        else:
            listNum = []
        return listNum


if __name__ == "__main__":
    app = QApplication(sys.argv)
    myWin = MyMainWindow()
    myWin.show()
    sys.exit(app.exec_())